import numpy as np
from collections import Counter
from sklearn.datasets import fetch_20newsgroups
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.model_selection import train_test_split


class NaiveBayes:
    def __init__(self, alpha=1.0):
        """
        初始化朴素贝叶斯分类器。

        参数:
            alpha (float): 拉普拉斯平滑参数，用于避免在计算概率时出现零概率问题，默认值为1.0。
        """
        self.alpha = alpha
        self.class_priors = {}  # 存储每个类别的先验概率 P(C)
        self.feature_likelihoods = {}  # 存储每个特征在不同类别下的条件概率 P(X|C)
        self.classes = []  # 存储所有可能的类别

    def fit(self, X, y):
        """
        根据提供的训练数据拟合朴素贝叶斯模型。

        参数:
            X (numpy.ndarray): 训练数据矩阵，其中每一行代表一个样本，每一列代表一个特征。
            y (numpy.ndarray): 训练数据标签向量，每个元素对应X中样本的类别。
        """
        self.classes = np.unique(y)  # 获取所有类别
        num_samples, num_features = X.shape

        # 计算每个类别的先验概率 P(C)，即该类别在所有样本中的比例
        self.class_priors = {c: np.mean(y == c) for c in self.classes}

        # 对每个类别计算条件概率 P(X|C)，应用拉普拉斯平滑以避免零概率问题
        self.feature_likelihoods = {}
        for c in self.classes:
            X_c = X[y == c]  # 获取属于类别c的所有样本

            likelihoods = {}
            for feature_idx in range(num_features):
                # 统计每个特征值在类别c中的出现次数
                feature_counts = Counter(X_c[:, feature_idx])

                # 计算每个特征的唯一值数量，并考虑未观察到的值（+1）
                unique_feature_values = len(feature_counts) + 1
                smoothed_counts = {feat_val: count + self.alpha for feat_val, count in feature_counts.items()}
                total_count = X_c.shape[0] + self.alpha * unique_feature_values

                # 将每个特征值的条件概率存储起来，并为未观察到的值设置默认概率
                likelihoods[feature_idx] = {feat_val: count / total_count for feat_val, count in smoothed_counts.items()}
                likelihoods[feature_idx][None] = self.alpha / total_count

            self.feature_likelihoods[c] = likelihoods

    def predict(self, X):
        """
        使用训练好的模型对新数据进行预测。

        参数:
            X (numpy.ndarray): 预测数据矩阵，结构与训练数据相同。

        返回:
            numpy.ndarray: 包含预测类别的数组。
        """
        predictions = []
        for sample in X:
            posteriors = {}
            for c in self.classes:
                # 初始化后验概率为类别c的先验概率的对数，防止数值下溢
                posterior = np.log(self.class_priors[c])

                # 累加每个特征的条件概率的对数值到后验概率上
                for feature_idx, feature_value in enumerate(sample):
                    likelihood = self.feature_likelihoods[c][feature_idx].get(feature_value, self.feature_likelihoods[c][feature_idx][None])
                    posterior += np.log(likelihood)

                posteriors[c] = posterior

            # 选择具有最高后验概率的类别作为预测结果
            predictions.append(max(posteriors, key=posteriors.get))

        return np.array(predictions)


def load_data():
    """
    加载并预处理20Newsgroups数据集，将文本转换为二进制特征向量，并划分训练集和测试集。

    返回:
        tuple: 包含训练集和测试集的数据和标签。
    """
    # 仅选择两个类别以简化问题
    newsgroups = fetch_20newsgroups(subset='all', categories=['sci.space', 'rec.autos'])

    # 文本向量化：使用CountVectorizer转换文本数据为特征向量
    vectorizer = CountVectorizer(binary=True, stop_words='english', max_features=1000)
    X = vectorizer.fit_transform(newsgroups.data).toarray()
    y = newsgroups.target

    # 数据集划分：80%训练集，20%测试集，保证结果可复现
    return train_test_split(X, y, test_size=0.2, random_state=42)


if __name__ == "__main__":
    # 加载数据集并获取训练集和测试集
    X_train, X_test, y_train, y_test = load_data()

    # 创建朴素贝叶斯分类器实例，并用训练集进行训练
    model = NaiveBayes(alpha=1.0)
    model.fit(X_train, y_train)

    # 对测试集进行预测
    y_pred = model.predict(X_test)

    # 计算并打印测试集上的准确率
    accuracy = np.mean(y_pred == y_test)
    print(f"测试集准确率: {accuracy:.4f}")